import sys
import os
import time
import datetime
import random

reload( sys )
sys.setdefaultencoding('utf-8')

class DLSSHUtil():

    tail = '\n\r'
    kill = 'DLKILL'
    def getSSHShell(self):
        from common import GetFixPath
        fix_path = GetFixPath()
        fix_path = os.path.abspath(fix_path)
        if sys.platform == 'win32':
            ssh_shell = fix_path+ '\\ext\\plink.exe '
        else:
            ssh_shell = fix_path+ '/ext/plink '
        return ssh_shell

    def getSCP(self):
        from common import GetFixPath
        fix_path = GetFixPath()
        fix_path = os.path.abspath(fix_path)
        if sys.platform == 'win32':
            ssh_shell = fix_path+'\\'+'ext\\pscp.exe '
        else:
            ssh_shell = fix_path+'/'+'ext/pscp '
        return ssh_shell

    def initConn(self,nodelist,isAdd=False):

        def EnableInterrupt(node):
            ostr = "stty isig intr '{0}' -echoctl; trap '/bin/true' SIGINT".format(self.kill)
            self.doCommand(node, ostr)
        def GetPtyNo(node):
            o = "w |awk '{print $2,$8}'|grep -w w|awk '{print $1}'"
            out = self.doCommand(node, o)
            return out

        def AddNewNode(loginStr,passwd):
            #print('add new node')
            out = ''
            ssh_shell = " echo y | "
            ssh_shell += self.getSSHShell()
            ssh_shell += ' -pw ' + passwd + ' '+loginStr
            ssh_shell += ' exit '
            print('get ssh key:' + loginStr)
            conn = Popen(ssh_shell, stdin=PIPE, stdout=PIPE, shell=True)
            while True:
                if conn.poll() == False:
                    break
                # out = Popen.recv_some(conn)
                # if out.find('error') > -1:
                    # conn.kill()
                    # conn = None
                    # break
            #out += 'ret code:'+str(conn.returncode)
            return out
        def login(loginStr,passwd):
            tail = self.tail
            shell = self.getSSHShell()
            shell += ' -pw ' + passwd + ' '+loginStr
            conn = Popen(shell, stdin=PIPE, stdout=PIPE, stderr=PIPE,shell=True)
            print('ssh connect:' + loginStr)
            while True:
                res = Popen.recv_some(conn)
                if res.find('password:') > 0:
                    Popen.send_all(conn, passwd + tail)
                elif res.find(']') > 0:
                    break
                # elif res.find('error')>-1 or res.find('ERROR') > -1:
                    # conn.kill()
                    # conn = None
                    # break
                #elif res.find('y/n') > 0:
                #    Popen.send_all(conn, tail)
            return conn

        conns = []
        tmpnodelist = []
        if isinstance(nodelist, dict):
            tmpnodelist = map(lambda x:[x,nodelist.get(x)], nodelist.keys())
        else:
            tmpnodelist = nodelist
        for (nodestr,passwd) in tmpnodelist:
            #try:
            if isAdd :
                out = AddNewNode(nodestr,passwd)
                #print(AddNewNode,out)
            conn = login(nodestr, passwd)
            conn_d = {'conn':conn, 'connStr':nodestr,'passwd':passwd,'kill':False}
            if conn is not None:
                conns.append(conn_d)

                #支持终止进程
                EnableInterrupt(conn_d)

                #第二种终止的方法
                ptsno = GetPtyNo(conn_d)
                conn_d['ptsno'] = ptsno
            # except Exception, ex:
                # print(ex, sys.exc_info()[:2])
                # continue

        return conns
    def closeConn(self,nodelist):
        tail = self.tail
        for node in nodelist:
            conn = node.get('conn')
            Popen.send_all(conn, 'exit' + tail)
            conn.wait()

    def killProcess(self, node):
        tmp_nodeList = self.initConn({node.get('connStr'):node.get('passwd')})

        ptsno = node.get('ptsno')
        o = " ps aux|grep -v 'grep\|sshd\|bash' |grep '"
        o += ptsno
        o += "'|awk '{print $2}'|sudo xargs kill"
        print(o)
        str = ""
        for node in tmp_nodeList:
            out = self.doCommand(node, o)
            print(out)
            str += out
        self.closeConn(tmp_nodeList)
        return str

    def killProcess2(self, node):
        out = self.doCommand(node, self.kill)
        out = self.doCommand(node, 'finish_doCommand')
        return out
    def putFile(self, node, filepath,destpath):
        #pscp -pw 235711 plink root@172.16.40.149:/opt/foo
        scp_shell = self.getSCP() + " -pw "
        scp_shell += node.get('passwd') + ' '
        scp_shell += filepath +' '
        scp_shell += node.get('connStr')+':'
        scp_shell += destpath
        putConn = Popen(scp_shell, stdin=PIPE, stdout=PIPE, shell=True)
        while True:
            if putConn.poll() == False:
                break

        out = 'ret code:'+str(putConn.returncode)
        return out

    def getFile(self, conn, srcpath, localpath):
        scp_shell = self.getSCP() + " -pw "
        scp_shell += conn.get('passwd') + ' '
        scp_shell += conn.get('connStr')+':'
        scp_shell += srcpath +' '
        scp_shell += localpath
        getConn = Popen(scp_shell, stdin=PIPE, stdout=PIPE, shell=True)
        while True:
            if getConn.poll() == False:
                break
        out = 'ret code:'+str(getConn.returncode)
        return out



    def doCommandAndWaitfor(self, node, cmd, response, t=.1, e=1, tr=5, stderr=0):
        tail = self.tail
        conn = node.get('conn')
        Popen.send_all(conn,cmd+tail)
        waitfor = '[.\r\n]*'+cmd+'[.\r\n]*'+response
        prog = re.compile(waitfor)
        print Popen.recv_some_restring(conn, prog, t, e, tr, stderr)

    def doCommandStr(self, node, commandStr,expect=None):
        tail = self.tail
        passwd = node.get('passwd')
        loginStr = node.get('connStr')
        shell = self.getSSHShell()
        shell += ' -pw ' + passwd + ' '+loginStr+ ' '+commandStr
        print(shell)
        conn = Popen(shell, stdin=PIPE, stdout=PIPE, shell=True)

        out = str(conn.poll())
#        res = Popen.recv_some(conn,e=0)
#        while True:
#            res = Popen.recv_some(conn,e=0)
#            out += res
#            if conn.returncode > -1:
#                break
#            elif len(res) > 0:
#                print(res)
#                isOpen = True
#            elif len(res) < 1 and isOpen:
#                break
        return out

    def doCommand(self, node, commandStr, expect=None):
        tail = '\n\r'
        endFlag = 'finish_doCommand'
        subCommand = ' ; echo ' +endFlag
        conn = node.get('conn')
        iskill = node.get('kill')
        if iskill != None and iskill:
            return "task is killed, command not send"

        totalCommand = commandStr + subCommand + tail
        Popen.send_all(conn, totalCommand)
        out = ''
        flagCount = 0
        while True:
            str = Popen.recv_some(conn, e=0)
            #print(str)
            if str.find(commandStr) > -1:
                str = str[len(totalCommand):]
                str = str.replace(commandStr,'')

            if str.find(subCommand) > -1:
                startPos = str.find(subCommand)+len(subCommand)
                str = str[startPos:]

            if expect is not None:
                res = expect.deal(str)
                if res is not None:
                    Popen.send_all(conn, res + tail)

            if str.find(endFlag) > -1:
                flagCount += 1
                endPos = str.find(endFlag)
                str = str[:endPos]
                out += str.replace('\r','').replace('\n','')
                if flagCount > 0:
                    break
            elif conn.returncode != None:
                #print(" return code :"+str(conn.returncode))
                out += str.replace('\r','').replace('\n','')
                break

            out += str
        return out



    def checkLinkOk(self, nodelist):
        nodeConns = self.initConn(nodelist, True)
        for node in nodeConns:
            out = self.doCommand(node,"ls -l")
            #print(out)
        self.closeConn(nodeConns)


    def sendPlink(self,nodelist):

        nodeConns = self.initConn(nodelist)
        for node in nodeConns:
            self.putFile(node, self.getSSHShell(),'/usr/bin/')
            self.putFile(node, self.getSCP(),'/usr/bin/')
            self.doCommand(node,"chmod 755 /usr/bin/plink")
            self.doCommand(node,"chmod 755 /usr/bin/pscp")
        self.closeConn(nodeConns)

    def getRemoteFile(self, nodelist, src, dest):
        nodeConns = self.initConn(nodelist)
        for node in nodeConns:
            self.getFile(node, src, dest+node.get("connStr"))

        self.closeConn(nodeConns)